import logging
from typing import List

from gpt.model.base import ModelInstance
from gpt.model.cluster.controller_base import ModeControllerBase
from gpt.model.cluster.registry_base import ModelRegistryBase

logger = logging.getLogger(__name__)


class LocalModelController(ModeControllerBase):
    def __init__(self, registry: ModelRegistryBase) -> None:
        self.registry = registry

    async def register_instance(self, instance: ModelInstance) -> bool:
        return await self.registry.register_instance(instance)

    async def deregister_instance(self, instance: ModelInstance) -> bool:
        return await self.registry.deregister_instance(instance)

    async def get_all_instances(
            self, model_name: str = None, healthy_only: bool = False
    ) -> List[ModelInstance]:
        logger.info(
            f"Get all instances with {model_name}, healthy_only: {healthy_only}"
        )
        if not model_name:
            return await self.registry.get_all_model_instances(
                healthy_only=healthy_only
            )
        else:
            return await self.registry.get_all_instances(model_name, healthy_only)

    async def send_heartbeat(self, instance: ModelInstance) -> bool:
        return await self.registry.send_heartbeat(instance)
